"""
!/usr/bin/env python
-*- coding: utf-8 -*-
@CreateTime : 2024/11/13 20:11
@Author  :    AnimateX
@Contact :    animatex@163.com
@File    :    rectify_rmsl.py
@License :    Copyright © 2024 AnimateX. All rights reserved.
@Version :    rectify_rmsl_2024/11/13.0.1

-------------------------------------------------------------------------------
# @Description:


    ---------------------------------------------------------------------------
    [Update History]:
        2024/11/13:

-------------------------------------------------------------------------------
"""
"""
!/usr/bin/env python
-*- coding: utf-8 -*-
@CreateTime : 2024/11/13 17:11
@Author  :    AnimateX
@Contact :    animatex@163.com
@File    :    rectify.py
@License :    Copyright © 2024 AnimateX. All rights reserved.
@Version :    rectify_2024/11/13.0.1

-------------------------------------------------------------------------------
# @Description:

    ---------------------------------------------------------------------------
    [Update History]:
        2024/11/13:

-------------------------------------------------------------------------------
"""
import os
import cv2
import numpy as np
from tqdm import tqdm
from natsort import natsorted


class RMSLRectify:
    def __init__(self,
                 yaml_path: str,
                 rotation_angle=90,
                 en_binning=False,
                 en_gamma=False,
                 gamma_val=2.2,
                 ir_datatype='u16',
                 en_demosaic=False,
                 demosaic_pattern='bggr'):
        self.yaml_path = yaml_path

        self.rotation_angle = rotation_angle
        self.enable_binning = en_binning
        self.enable_gamma = en_gamma
        self.gamma_value = gamma_val
        self.output_depth = ir_datatype
        self.enable_demosaic = en_demosaic
        self.demosaic_pattern = demosaic_pattern
        self.parser_config_info_from_yaml()

    def parser_config_info_from_yaml(self):
        if not os.path.exists(self.yaml_path):
            raise FileNotFoundError(f' [Error] File {self.yaml_path} not found.')
        fs = cv2.FileStorage(self.yaml_path, cv2.FILE_STORAGE_READ)

        # Image dimensions
        self.left_width = int(fs.getNode("l_width").real())
        self.left_height = int(fs.getNode("l_height").real())
        self.right_width = int(fs.getNode("r_width").real())
        self.right_height = int(fs.getNode("r_height").real())

        # Camera intrinsics and distortion coefficients
        self.K_l = fs.getNode("K_l").mat()
        self.D_l = fs.getNode("D_l").mat()
        self.K_r = fs.getNode("K_r").mat()
        self.D_r = fs.getNode("D_r").mat()

        # Stereo rectification parameters
        self.R1 = fs.getNode("R1").mat()
        self.R2 = fs.getNode("R2").mat()
        self.P1 = fs.getNode("P1").mat()
        self.P2 = fs.getNode("P2").mat()
        self.Q = fs.getNode("Q").mat()

        # Valid region
        # self.validRoi0 = fs.getNode("validRoi0").mat().astype(int)
        # self.validRoi1 = fs.getNode("validRoi1").mat().astype(int)

        fs.release()

    @staticmethod
    def read_16bit_image(image_path):
        """
        Read 16-bit PNG image
        """
        img = cv2.imread(str(image_path), cv2.IMREAD_UNCHANGED)
        if img is None:
            raise ValueError(f"Cannot read image: {image_path}")
        if img.dtype != np.uint16:
            raise ValueError(f"Image is not 16-bit: {image_path}")
        return img

    @staticmethod
    def rotate_image(img, angle):
        """
        Rotate image using flip and transpose (only supports 0°, 90°, 180°, 270°)
        Args:
            img: Input image
            angle: Rotation angle (must be a multiple of 90)
        Returns:
            Rotated image
        """
        # Ensure angle is a multiple of 90
        angle = angle % 360
        if angle % 90 != 0:
            raise ValueError("Rotation angle must be a multiple of 90")

        # Normalize to 0-270 degrees
        angle = angle % 360

        if angle == 0:
            return img
        elif angle == 90:
            # First transpose then flip horizontally
            return cv2.flip(cv2.transpose(img), 1)
        elif angle == 180:
            # Flip vertically and horizontally
            return cv2.flip(img, -1)
        elif angle == 270:
            # First transpose then flip vertically
            return cv2.flip(cv2.transpose(img), 0)

    def demosaic_image(self, img):
        """
        Demosaic the image and convert to grayscale
        """
        if self.demosaic_pattern not in ['bggr', 'rggb', 'gbrg', 'grbg']:
            raise ValueError(f"Unsupported demosaic pattern: {self.demosaic_pattern}")

        # Demosaic
        if self.demosaic_pattern == 'bggr':
            demosaiced = cv2.cvtColor(img, cv2.COLOR_BayerBG2RGB)
        elif self.demosaic_pattern == 'rggb':
            demosaiced = cv2.cvtColor(img, cv2.COLOR_BayerRG2RGB)
        elif self.demosaic_pattern == 'gbrg':
            demosaiced = cv2.cvtColor(img, cv2.COLOR_BayerGB2RGB)
        else:  # 'grbg'
            demosaiced = cv2.cvtColor(img, cv2.COLOR_BayerGR2RGB)
            # Save demosaiced image for debugging

        # debug_filename = f"debug_demosaic.png"
        # cv2.imwrite(debug_filename, (demosaiced/4).astype(np.uint8))
        # Convert to grayscale
        return cv2.cvtColor(demosaiced, cv2.COLOR_RGB2GRAY)

    def preprocess_image(self, img):
        """
        Image preprocessing: demosaic, rotation, binning, and gamma correction
        """
        # Demosaic
        if self.enable_demosaic:
            img = self.demosaic_image(img)

        # Rotate image (only supports 0°, 90°, 180°, 270°)
        if self.rotation_angle != 0:
            img = self.rotate_image(img, self.rotation_angle)

        # Binning processing
        if self.enable_binning:
            img = cv2.resize(img, None, fx=0.5, fy=0.5, interpolation=cv2.INTER_AREA)

        # Gamma correction
        if self.enable_gamma:
            # Normalize 16-bit image to 0-1
            img_norm = img.astype(float) * 64 / 65535.0
            # Apply gamma correction
            img_gamma = np.power(img_norm, 1 / self.gamma_value)

            # Normalize
            img_gamma_min = img_gamma.min()
            img_gamma_max = img_gamma.max()

            # Output 8-bit or 16-bit based on configuration
            if self.output_depth == 'u8':
                # img = (img_gamma * 255.0).astype(np.uint8)
                img = ((img_gamma - img_gamma_min) / (img_gamma_max - img_gamma_min) * 255).astype(np.uint8)
            else:
                # img = (img_gamma * 65535.0).astype(np.uint16)
                img = ((img_gamma - img_gamma_min) / (img_gamma_max - img_gamma_min) * 65535).astype(np.uint8)

        return img

    def rectify_images(self, left_img, right_img):
        """
        Perform epipolar rectification on stereo images using pre-calculated rectification parameters
        """
        # Calculate reprojection mappings
        left_map1, left_map2 = cv2.initUndistortRectifyMap(
            self.K_l,
            self.D_l,
            self.R1,
            self.P1,
            (self.left_width, self.left_height),
            cv2.CV_16SC2,
        )

        right_map1, right_map2 = cv2.initUndistortRectifyMap(
            self.K_r,
            self.D_r,
            self.R2,
            self.P2,
            (self.right_width, self.right_height),
            cv2.CV_16SC2,
        )

        # Perform rectification
        rect_left = cv2.remap(left_img, left_map1, left_map2, cv2.INTER_LINEAR)
        rect_right = cv2.remap(right_img, right_map1, right_map2, cv2.INTER_LINEAR)

        # Crop to valid region (optional)
        # if self.validRoi0 is not None and self.validRoi1 is not None:
        #     x, y, w, h = self.validRoi0
        #     rect_left = rect_left[y : y + h, x : x + w]
        #     x, y, w, h = self.validRoi1
        #     rect_right = rect_right[y : y + h, x : x + w]

        return rect_left, rect_right

    def process_stereo_pair(self, left_path, right_path):
        """
        Process a single pair of stereo images
        """
        try:
            # Read images
            left_img = RMSLRectify.read_16bit_image(left_path)
            right_img = RMSLRectify.read_16bit_image(right_path)

            # Preprocess
            left_img = self.preprocess_image(left_img)
            right_img = self.preprocess_image(right_img)

            # Epipolar rectification
            rect_left, rect_right = self.rectify_images(left_img, right_img)

            # save
            left_output = left_path.replace('_left', '_rectify_speckle-left')
            right_output = right_path.replace('_right', '_rectify_speckle-right')
            cv2.imwrite(left_output, rect_left)
            cv2.imwrite(right_output, rect_right)

            return True
        except Exception as e:
            print(f"Error processing image: {left_path}: {str(e)}")
            return False


def get_image_path_list(test_dir: str, key='_left.png') -> list:
    image_paths = []
    with os.scandir(test_dir) as entries:
        for entry in entries:
            if entry.is_file() and key in entry.name:
                image_paths.append(entry.path)
            elif entry.is_dir():
                image_paths.extend(get_image_path_list(entry.path))

    # normal path format.
    image_paths = natsorted([img_path.replace('\\', '/') for img_path in image_paths])

    return image_paths


def generate_rectify_speckle_images(test_dir: str, rmsl_yaml_path: str, key='_left.png'):
    if not os.path.exists(test_dir):
        raise FileNotFoundError(test_dir)

    if not os.path.exists(rmsl_yaml_path):
        raise FileNotFoundError(rmsl_yaml_path)

    rmsl_rectifier = RMSLRectify(rmsl_yaml_path)

    left_image_path_list = get_image_path_list(test_dir, key)
    for idx, left_img_path in enumerate(tqdm(left_image_path_list, position=0, desc='Rectify...', colour='cyan')):
        right_img_path = left_img_path.replace('_left.png', '_right.png')
        if not os.path.exists(right_img_path):
            print(f' [Warning] Rectify speckle images not found for {right_img_path}')
            continue

        rmsl_rectifier.process_stereo_pair(left_img_path, right_img_path)


if __name__ == '__main__':
    rmsl_test_dataset_path = '/home/dataset_b/rmslTestDataset/rockchip/rmsl321/v241111'
    rmsl_yaml_cfg_path = '/home/dataset_b/rmslTestDataset/rockchip/rmsl321/v241111/calib_param.yaml'
    generate_rectify_speckle_images(test_dir=rmsl_test_dataset_path,
                                    rmsl_yaml_path=rmsl_yaml_cfg_path,
                                    key='_left.png')
